import os
import cv2
import numpy as np
from torch.jit import annotate
def vis_parsing_maps(im, parsing_anno, save_im=False, save_path='',im_name='1.png'):
    stride=0
    # Colors for all 20 parts
    part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0],
                   [255, 0, 85], [255, 0, 170],
                   [0, 255, 0], [85, 255, 0], [170, 255, 0],
                   [0, 255, 85], [0, 255, 170],
                   [0, 0, 255], [85, 0, 255], [170, 0, 255],
                   [0, 85, 255], [0, 170, 255],
                   [255, 255, 0], [255, 255, 85], [255, 255, 170],
                   [255, 0, 255], [255, 85, 255], [255, 170, 255],
                   [0, 255, 255], [85, 255, 255], [170, 255, 255]]

    im = np.array(im)
    vis_im = im.copy().astype(np.uint8)
    vis_parsing_anno = parsing_anno.copy().astype(np.uint8)
    print(np.unique(vis_parsing_anno), vis_parsing_anno.shape)
    # vis_parsing_anno = cv2.resize(vis_parsing_anno, None, fx=stride, fy=stride, interpolation=cv2.INTER_NEAREST)
    vis_parsing_anno_color = np.zeros((vis_parsing_anno.shape[0], vis_parsing_anno.shape[1], 3)) + 255

    num_of_class = np.max(vis_parsing_anno)

    for pi in range(1, num_of_class + 1):
        index = np.where(vis_parsing_anno == pi)
        vis_parsing_anno_color[index[0], index[1], :] = part_colors[pi]

    vis_parsing_anno_color = vis_parsing_anno_color.astype(np.uint8)
    print(vis_parsing_anno_color.shape, vis_im.shape)

    vis_im = cv2.addWeighted(cv2.cvtColor(vis_im, cv2.COLOR_RGB2BGR), 0.4, vis_parsing_anno_color, 0.6, 0)

    # Save result or not
    if save_im:
        cv2.imwrite(save_path+'/anno/'+im_name[:-4] +'.png', vis_parsing_anno)
        cv2.imwrite(save_path+'/weights_img/'+im_name, vis_im, [int(cv2.IMWRITE_JPEG_QUALITY), 100])

root_dir = './data/cvpr/train/'
data_dir = root_dir+'image/'
ann_dir = root_dir+'/seg/'
for id in os.listdir(data_dir):
    for img in os.listdir(data_dir+id):
        im = cv2.imread(data_dir+id+'/'+img)
        anno = cv2.imread(ann_dir+id+'/'+img.replace('jpg', 'png'), cv2.IMREAD_GRAYSCALE)
        print(im,anno)
        vis_parsing_maps(im, anno, save_im=True, save_path=root_dir+'vis', im_name=img)